import os

from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate
from pydantic import BaseModel, Field

from domain.model.register import CHAT_MODEL_MANAGER
from infrastructure.code_extractor import GoCodeExtractor
from infrastructure.utils.env import set_langsmith
from infrastructure.utils.process import timer_context

system_message_content = """角色：铜雀台活动开发者
背景：活动开发需要将 Excel 配置项定义为结构体，现在已经将配置项转化为 JSON 二维数组的格式。用户会给你提供 JSON 格式的配置项和可复用的全家桶数据。
职责：帮助用户判断提供的配置项（JSON 二维数组）是否可以复用全家桶结构体，要求返回结构化结果。

数据定义：
+ 全家桶以 JSON 数组的格式提供
+ 每个元素是一个可复用的全家桶对象
+ 全家桶对象格式定义：
+ key 代表全家桶结构体名称，全家桶结构体又分为铜雀台（以 Tqt 为前缀）和服务端（以 Svr 为前缀），value 为全家桶的具体定义。
+ 如果 key 为 `table` 表示该配置项的 JSON 格式，你可以利用 `table` 字段判断与用户传入的配置项是否高度匹配，是否可复用。
+ 示例：
```json
{
  "SvrRank": "type SvrRank struct {\n\tRankFrom int        `json:\"rank_from\"`\n\tRankTo   int        `json:\"rank_to\"`\n\tAward    []SvrAward `json:\"award\"`\n}",
  "TqtRank": "type TqtRank struct {\n\tRankFrom int        `json:\"rankFrom\" header:\"开始名次\" validate:\"gt=0\"`\n\tRankTo   int        `json:\"rankTo\" header:\"结束名次\" validate:\"gte=-1\"`\n\tAward    []TqtAward `json:\"award\" header:\"奖励内容\" validate:\"gt=0,dive\"`\n}",
  "table": [
    [
      "个人排名奖励"
    ],
    [
      "开始名次",
      "结束名次",
      "奖励内容",
      "奖励内容",
      "奖励内容"
    ],
    [
      1,
      20,
      "金币*1",
      "金币*1",
      "金币*1"
    ],
    [
      21,
      99,
      "金币*1",
      "金币*1",
      "金币*1"
    ],
    [
      99,
      -1,
      "金币*1",
      "金币*1",
      "金币*1"
    ]
  ]
}
```

输出要求：
+ 禁止任何自然语言解释
+ 禁止无关的文字输出，需要遵循定义以下定义的格式：
{{format_instructions}}
"""

value_path = '../../data/code/value'


class Result(BaseModel):
    matched: bool = Field(..., description="是否可复用")
    tqt_structs: list[str] = Field(default=[], description="铜雀台结构体具体定义")
    svr_structs: list[str] = Field(default=[], description="服务端结构体具体定义")
    item_name: str = Field(..., description="配置项名称")
    table: str = Field(..., description="配置项JSON二维数组")


def open_go_file(file_name: str) -> str:
    """
    打开go文件
    :param file_name: 文件名称
    :return: 文件内容
    """
    with open(os.path.join(value_path, file_name), 'r', encoding='utf-8') as f:
        return f.read()


def construct_family_bucket_chain(is_debug: bool = False):
    """
    构造全家桶 Runnable链
    需传入 all_config_item, config_item
    """
    if is_debug:
        set_langsmith()
    chat_model = CHAT_MODEL_MANAGER.get_model_instance_by_name('glm-4-plus')
    parser = JsonOutputParser(pydantic_object=Result)
    prompt = ChatPromptTemplate(
        messages=[
            ('system', system_message_content),
            ('human', '提供给你的全家桶数据:\n{{all_config_item}}'),
            ('human', '判断该配置项是否可复用:\n{{config_item}}')
        ],
        template_format="jinja2",
        input_variables=['all_config_item', 'config_item'],
        partial_variables={"format_instructions": parser.get_format_instructions()},
    )
    return prompt | chat_model | parser


def scan_all_family_bucket_to_config_item():
    """扫描全家桶并转化为配置"""
    all_item = []
    go_file_paths = [file_name for file_name in os.listdir(value_path) if file_name.endswith('.go')]
    for path in go_file_paths:
        go_code = open_go_file(path)
        all_item.append(GoCodeExtractor.extract_from_content(go_code))
    return [item for item in all_item if item]


def family_bucket_tool(
        config_item: list[list] = Field(..., description="配置项的JSON二维数组")
) -> dict:
    with timer_context('扫描全家桶'):
        all_config_item = scan_all_family_bucket_to_config_item()
    chain = construct_family_bucket_chain(True)
    with timer_context('大模型推理'):
        return chain.invoke({
            'all_config_item': all_config_item,
            'config_item': config_item
        })


if __name__ == '__main__':
    with timer_context('全家桶工具'):
        result = family_bucket_tool([
            ["礼包配置"],
            ["礼包类型", "编号", "礼包名称", "货币种类", "货币数量", "购买项id", "价格", "第几天开始", "第几天结束",
             "开始时间", "结束时间", "价值千分比", "限购次数", "是否每日重置"],
            ["无极礼包", 41, "无极礼包1", "金币", 100, "", "", 1, 2, "00:00:00", "23:59:59", 200, 100, "是"],
            ["无极礼包", 42, "无极礼包2", "金币", 50, "", "", 3, 4, "00:00:00", "23:59:59", 90, 100, "否"]
        ])
        print(result)
    with timer_context('调用gemini'):
        chat_model = CHAT_MODEL_MANAGER.get_model_instance_by_name('gemini-2.5-flash-preview-05-20')
        print(chat_model.invoke('模型根据用户提问检索已存在知识库得到参考，以实现RAG流程的实现思路。'))
